package com.yth.utils;


import lombok.EqualsAndHashCode;
import lombok.Getter;
import lombok.ToString;

import java.lang.ref.ReferenceQueue;
import java.lang.ref.WeakReference;
import java.util.Collections;
import java.util.LinkedHashMap;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.LongAdder;
import java.util.function.Function;


/**
 * 一个使用弱引用来存储值的Map实现。
 * 当值对象不再被外部强引用时，这些Entries会被自动移除。
 *
 * @param <K> 键的类型
 * @param <V> 值的类型
 * @author yutianhong
 * @version 1.0
 * @since 2025/3/25 17:21
 */
public class WeakValueLoadingCache<K, V> {

    @Getter
    public static class CacheLoadingException extends RuntimeException {
        private final Object key;

        CacheLoadingException(Object key, Throwable cause) {
            super("Error loading value for key: " + key, cause);
            this.key = key;
        }
    }

    /**
     * Statistics about the performance of a {@link WeakValueLoadingCache}. Instances of this class are immutable.
     */
    @Getter
    @EqualsAndHashCode
    @ToString
    public static final class CacheStats {
        private final long hitCount;
        private final long missCount;

        CacheStats(long hitCount, long missCount) {
            this.hitCount = hitCount;
            this.missCount = missCount;
        }
    }

    private final Map<K, WeakValue> map;
    private final ReferenceQueue<V> queue = new ReferenceQueue<>();
    private final Function<K, V> loader;

    private final LongAdder hitCount = new LongAdder();
    private final LongAdder missCount = new LongAdder();

    private class WeakValue extends WeakReference<V> {
        private final K key;

        WeakValue(V value, K key, ReferenceQueue<V> queue) {
            super(value, queue);
            this.key = key;
        }
    }

    public WeakValueLoadingCache(Function<K, V> loader, int initialCapacity) {
        Objects.requireNonNull(loader);
        this.loader = loader;
        this.map = new ConcurrentHashMap<>(initialCapacity);
    }

    /**
     * Returns the value associated with {@code key} in this cache, or {@code null} if there is no
     * cached value for {@code key}.
     */
    public V getIfPresent(K key) {
        Objects.requireNonNull(key, "key cannot be null");
        processQueue();
        WeakValue ref = map.get(key);
        return ref != null ? ref.get() : null;
    }

    /**
     * Returns the value associated with {@code key} in this cache, obtaining that value from {@code
     * loader} if necessary.
     *
     * @throws CacheLoadingException if a checked exception was thrown while loading the value
     */
    public V get(K key) throws CacheLoadingException {
        Objects.requireNonNull(key, "key cannot be null");
        processQueue();
        Object[] loadedValueRef = new Object[1];
        Function<? super K, ? extends WeakValue> mappingFunction =
                (K k) -> {
                    V value;
                    try {
                        value = loader.apply(k);
                    } catch (Throwable e) {
                        throw new CacheLoadingException(k, e);
                    }
                    if (value == null) {
                        throw new CacheLoadingException(k, new IllegalArgumentException("Loader returned null for key: " + k));
                    }
                    loadedValueRef[0] = value;      // 持有value引用，防止极端场景下，在compute方法执行期间发生gc
                    return new WeakValue(value, k, queue);
                };
        WeakValue ref = map.computeIfAbsent(key, mappingFunction);
        if (loadedValueRef[0] == null) {
            hitCount.increment();
        } else {
            missCount.increment();
        }
        return ref.get();
    }

    /**
     * Returns a map of the values associated with {@code keys} in this cache. The returned map will
     * only contain entries which are already present in the cache.
     */
    public Map<K, V> getAllPresent(Iterable<? extends K> keys) {
        Map<K, V> result = new LinkedHashMap<>();
        for (K key : keys) {
            if (key != null && !result.containsKey(key)) {
                V value = getIfPresent(key);
                if (value != null) {
                    result.put(key, value);
                }
            }
        }
        return Collections.unmodifiableMap(result);
    }

    /**
     * Associates {@code value} with {@code key} in this cache. If the cache previously contained a
     * value associated with {@code key}, the old value is replaced by {@code value}.
     */
    public void put(K key, V value) {
        Objects.requireNonNull(key, "key cannot be null");
        Objects.requireNonNull(value, "value cannot be null");
        processQueue();
        map.put(key, new WeakValue(value, key, queue));
    }

    /**
     * Copies all the mappings from the specified map to the cache. The effect of this call is
     * equivalent to that of calling {@code put(k, v)} on this map once for each mapping from key
     * {@code k} to value {@code v} in the specified map. The behavior of this operation is undefined
     * if the specified map is modified while the operation is in progress.
     */
    public void putAll(Map<? extends K, ? extends V> m) {
        for (Map.Entry<? extends K, ? extends V> entry : m.entrySet()) {
            put(entry.getKey(), entry.getValue());
        }
    }

    /**
     * Discards any cached value for key {@code key}.
     */
    public void invalidate(K key) {
        Objects.requireNonNull(key, "key cannot be null");
        processQueue();
        map.remove(key);
    }

    /**
     * Discards any cached values for keys {@code keys}.
     */
    public void invalidateAll(Iterable<? extends K> keys) {
        for (K key : keys) {
            invalidate(key);
        }
    }

    /**
     * Discards all entries in the cache.
     */
    public void invalidateAll() {
        processQueue();
        map.clear();
    }

    /**
     * Performs any pending maintenance operations needed by the cache.
     */
    public void cleanup() {
        processQueue();
    }

    /**
     * Returns the approximate number of entries in this cache.
     */
    public int size() {
        processQueue();
        return map.size();
    }

    public boolean isEmpty() {
        processQueue();
        return map.isEmpty();
    }

    /**
     * Returns a current snapshot of this cache's cumulative statistics, or a set of default values if
     * the cache is not recording statistics. All statistics begin at zero and never decrease over the
     * lifetime of the cache.
     */
    public CacheStats stats() {
        return new CacheStats(hitCount.sum(), missCount.sum());
    }

    @SuppressWarnings("unchecked")
    private void processQueue() {
        WeakValue wv;
        while ((wv = (WeakValue) queue.poll()) != null) {
            map.remove(wv.key, wv);
        }
    }
}
